import pygame
import sys
from collections import deque

# Define colors
WALKABLE = (255, 182, 193)  # Pink for walkable paths
WALL = (0, 0, 255)  # Blue for walls
START = (0, 124, 255)  # Blue for start
EXIT = (255, 0, 0)  # Red for exit
VISITED = (128, 0, 128)  # Purple for visited cells
EXPLORED = (173, 216, 230)  # Light blue for cells during exploration
PATH = (0, 255, 0)  # Green for final path

# Initialize pygame
pygame.init()

# Set screen size
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("广度搜索法 (BFS)")

# Set font for displaying queue numbers
font = pygame.font.SysFont(None, 24)

# Calculate cell size based on maze size
def calculate_cell_size(maze):
    cell_width = WIDTH // len(maze[0])
    cell_height = HEIGHT // len(maze)
    return cell_width, cell_height

# Initialize queue number counter
queue_number = 0
queue_numbers = {}  # Dictionary to store queue numbers for each cell

# Store steps for next and previous functionality
bfs_steps = []  # This will store all the steps (state at each step)
current_step = -1  # Index to track the current step in the exploration

# BFS variables for interactive step-by-step control
queue = deque()
visited = set()
parent = {}
paused = True  # Start in paused mode
bfs_done = False  # Track if BFS is completed

def load_maze_from_file(filepath):
    """Load maze from the specified file."""
    with open(filepath, 'r') as f:
        maze = [list(line.strip()) for line in f.readlines()]
    return maze

def draw_maze(maze, cell_width, cell_height):
    screen.fill((0, 0, 0))  # Clear screen
    for i in range(len(maze)):
        for j in range(len(maze[0])):
            color = WALKABLE if maze[i][j] == '.' else WALL
            pygame.draw.rect(screen, color, (j * cell_width, i * cell_height, cell_width, cell_height))
            if maze[i][j] == 'S':
                pygame.draw.rect(screen, START, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'E':
                pygame.draw.rect(screen, EXIT, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == '#':
                pygame.draw.rect(screen, VISITED, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'P':
                pygame.draw.rect(screen, PATH, (j * cell_width, i * cell_height, cell_width, cell_height))
            elif maze[i][j] == 'X':
                pygame.draw.rect(screen, EXPLORED, (j * cell_width, i * cell_height, cell_width, cell_height))
            
            # Display queue number on explored cells
            if (i, j) in queue_numbers:
                text = font.render(str(queue_numbers[(i, j)]), True, (0, 0, 0))
                screen.blit(text, (j * cell_width + cell_width // 4, i * cell_height + cell_height // 4))
                
    pygame.display.flip()

def bfs_step(maze, cell_width, cell_height):
    """Perform a single step of BFS."""
    global queue_number, bfs_done

    if not queue or bfs_done:
        return

    x, y = queue.popleft()

    if maze[x][y] == 'E':
        # Backtrack to draw the final path
        draw_final_path(maze, parent, x, y, cell_width, cell_height)
        bfs_done = True
        return
    
    # Mark as currently being explored
    if maze[x][y] not in ['S', 'E']:
        maze[x][y] = 'X'  # Use 'X' for exploring cells
    
    # Save the current state for the next/previous steps
    save_current_state(maze)
    draw_maze(maze, cell_width, cell_height)

    # Explore neighbors
    for dx, dy in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
        nx, ny = x + dx, y + dy
        if 0 <= nx < len(maze) and 0 <= ny < len(maze[0]) and (nx, ny) not in visited and maze[nx][ny] != '*':
            queue.append((nx, ny))
            visited.add((nx, ny))
            parent[(nx, ny)] = (x, y)
            
            # Increment queue number and assign it to the cell
            queue_number += 1
            queue_numbers[(nx, ny)] = queue_number
    
    # After exploring all neighbors of this node, mark it as visited
    if maze[x][y] == 'X':
        maze[x][y] = '#'


def bfs_init(maze, x, y):
    """Initialize BFS."""
    global queue_number, queue, visited, parent, paused, bfs_done
    queue = deque([(x, y)])
    visited = set([(x, y)])
    parent = {(x, y): None}
    queue_number = 1
    queue_numbers[(x, y)] = queue_number  # Assign queue number to start node
    bfs_done = False
    paused = True  # Start in paused mode
    save_current_state(maze)
    draw_maze(maze, *calculate_cell_size(maze))


def draw_final_path(maze, parent, x, y, cell_width, cell_height):
    """Draw the final path once the exit is found."""
    while parent[(x, y)] is not None:
        maze[x][y] = 'P'  # Mark final path with 'P'
        x, y = parent[(x, y)]
        save_current_state(maze)
        draw_maze(maze, cell_width, cell_height)

def save_current_state(maze):
    global bfs_steps, current_step
    # Store the current state of the maze, queue numbers, and any other variables
    state = {
        'maze': [row[:] for row in maze],  # Copy the maze
        'queue_numbers': dict(queue_numbers)  # Copy queue numbers
    }
    bfs_steps.append(state)  # Add state to bfs_steps list
    current_step += 1  # Increment step counter

def load_state(step):
    """Load a specific state in the exploration."""
    global maze, queue_numbers, current_step
    # Load the state of the maze from the given step
    if 0 <= step < len(bfs_steps):
        state = bfs_steps[step]
        maze = [row[:] for row in state['maze']]  # Copy back the maze state
        queue_numbers = dict(state['queue_numbers'])  # Copy back queue numbers
        current_step = step
        draw_maze(maze, *calculate_cell_size(maze))

def main():
    global paused, current_step

    # Check if the user provided a maze file path
    if len(sys.argv) < 2:
        print("Usage: python bfs_visualization.py <maze_file.txt>")
        sys.exit(1)

    # Load the maze from the file
    maze_file = sys.argv[1]
    maze = load_maze_from_file(maze_file)

    # Calculate cell sizes based on the loaded maze
    cell_width, cell_height = calculate_cell_size(maze)

    # Find start position
    start_x, start_y = None, None
    for i in range(len(maze)):
        for j in range(len(maze[0])):
            if maze[i][j] == 'S':
                start_x, start_y = i, j

    if start_x is not None and start_y is not None:
        bfs_init(maze, start_x, start_y)

    clock = pygame.time.Clock()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    paused = not paused  # Pause or resume
                elif event.key == pygame.K_RIGHT:
                    if paused and current_step + 1 < len(bfs_steps):
                        load_state(current_step + 1)  # Move to the next step
                elif event.key == pygame.K_LEFT:
                    if paused and current_step - 1 >= 0:
                        load_state(current_step - 1)  # Move to the previous step

        if not paused:
            bfs_step(maze, cell_width, cell_height)  # Execute one step of BFS when not paused

        clock.tick(60)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()